import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from scipy.spatial import ConvexHull
import time
from utils.tools import read_json
import pickle
from tqdm import tqdm
import random
np.random.seed(int(time.time()))
random.seed(int(time.time()))

def dominates(a, b):
    return np.all(a <= b) and np.any(a < b)

class AMOPSO_ESE:
    def __init__(self, I, J, K, c, d, t, mu, C, pop_size=50, max_iter=100, archive_size=50):
        self.I = I  # 供应点数量
        self.J = J  # 受灾点数量
        self.K = K  # 物资种类
        
        # 强制转换为NumPy数组
        self.c = np.array(c, dtype=int)    # 供应量 (I, K)
        self.d = np.array(d, dtype=int)    # 需求量 (J, K)
        self.t = np.array(t, dtype=float)  # 运输时间 (I, J)
        self.mu = np.array(mu, dtype=float)  # 优先系数 (J,)
        
        self.C = C    # 车辆容量
        self.pop_size = pop_size
        self.max_iter = max_iter
        self.archive_size = archive_size
        
        # PSO参数
        self.omega = 0.9
        self.c1 = 1.424
        self.c2 = 1.424
        self.T = 10  # 状态切换阈值
        
        # 初始化种群
        self.particles = [self.initialize_particle() for _ in range(pop_size)]
        self.velocities = [np.zeros_like(p, dtype=float) for p in self.particles]
        self.pbest_pos = [p.copy() for p in self.particles]
        self.pbest_fit = []
        self.pbest_viol = []
        for p in self.particles:
            fit, viol = self.evaluate(p)
            self.pbest_fit.append(fit)
            self.pbest_viol.append(viol)
        
        # 外部档案和进化状态
        self.archive = []
        self.Hmin = float('inf')
        self.ES = 'EXPLOITATION'
        self.Counter = 0

    def initialize_particle(self):
        particle = np.zeros((self.I, self.J, self.K), dtype=int)
        i_list = list(range(self.I))
        k_list = list(range(self.K))
        random.shuffle(i_list)
        random.shuffle(k_list)
        for i in i_list:
            for k in k_list:
                remaining = self.c[i, k]
                j_list = list(range(self.J))
                random.shuffle(j_list)
                for j in j_list:
                    allocated_before = np.sum(particle[:i+1, j, k])
                    remaining_demand = self.d[j, k] - allocated_before
                    if remaining <= 0 or remaining_demand <= 0:
                        continue
                    max_assign = min(remaining, remaining_demand)
                    assign = np.random.randint(0, max_assign + 1)
                    particle[i, j, k] = assign
                    remaining -= assign
                # 处理剩余未分配量（确保供应总量守恒）
                current_total = np.sum(particle[i, :, k])
                if current_total < self.c[i, k]:
                    remaining = self.c[i, k] - current_total
                    j_list = np.random.permutation(self.J)
                    for j in j_list:
                        allocated_before = np.sum(particle[:, j, k])
                        remaining_demand = self.d[j, k] - allocated_before
                        if remaining_demand <= 0:
                            continue
                        max_add = min(remaining, remaining_demand)
                        particle[i, j, k] += max_add
                        remaining -= max_add
                        if remaining == 0:
                            break
        return particle

    def evaluate(self, particle):
        # 时间周转量计算（使用NumPy广播）
        
        f1 = np.max((particle > 0).astype(int) * self.t[:, :, np.newaxis])
        # tmp = np.max((particle > 0).astype(int) * self.t[:, :, np.newaxis], axis=(1,2)).tolist()
        # f3 = -1
        # for x in tmp:
        #     if x > 0 and (f3 == -1 or x < f3):
        #         f3 = x
        
        # 公平性计算
        received = np.sum(particle, axis=(0,2))  # 各受灾点接收总量 (J,)
        total_demand = np.sum(self.d, axis=1)    # (J,)
        unmet_ratio = 1 - received / (total_demand + 1e-9)
        unmet_ratio = np.clip(unmet_ratio, 0.0, 1.0)
        f2 = np.sum(self.mu * unmet_ratio)
        
        # # 空载率计算
        # total_weight = np.sum(particle, axis=2)  # (I,J)
        # a_ij = total_weight % self.C
        # non_zero = np.count_nonzero(a_ij)
        # sum_a = np.sum(a_ij)
        # f3 = 1 - sum_a/(self.C*non_zero) if non_zero > 0 else 0.0
        
        # 约束违反计算
        violation = 0
        # 需求约束
        received_per_demand = np.sum(particle, axis=0)  # (J,K)
        violation += np.sum(np.maximum(received_per_demand - self.d, 0))
        # 供应约束
        supplied_per_source = np.sum(particle, axis=1)  # (I,K)
        violation += np.sum(np.abs(supplied_per_source - self.c))
        
        # return np.array([f1, f2, f3]), violation
        return np.array([f1, f2]), 0

    def update_archive(self, new_solutions):
        for sol in new_solutions:
            is_dominated = False
            to_remove = []
            for idx, a in enumerate(self.archive):
                a_fit = a['fitness']
                s_fit = sol['fitness']
                if dominates(a_fit, s_fit):
                    is_dominated = True
                    break
                if dominates(s_fit, a_fit):
                    to_remove.append(idx)
            if not is_dominated:
                self.archive = [a for i, a in enumerate(self.archive) if i not in to_remove]
                self.archive.append(sol)
        
        # 拥挤距离排序保留
        if len(self.archive) > self.archive_size:
            fits = np.array([s['fitness'] for s in self.archive])
            # 简单按第一个目标排序后截断
            sorted_indices = np.argsort(fits[:, 0])
            self.archive = [self.archive[i] for i in sorted_indices[:self.archive_size]]

    def evolutionary_state_evaluation(self):
        if not self.archive:
            return
        # 计算到原点的距离
        AD = [np.linalg.norm(sol['fitness']) for sol in self.archive]
        Amin = min(AD)
        
        if self.ES == 'EXPLOITATION':
            if self.Hmin <= Amin:
                self.Counter += 1
                if self.Counter > self.T:
                    self.ES = 'EXPLORATION'
                    self.Counter = 0
            else:
                self.Hmin = Amin
                self.Counter = 0
        else:
            if Amin < self.Hmin:
                self.ES = 'EXPLOITATION'
                self.Hmin = Amin
                self.Counter = 0     

    def run(self):
        for _ in tqdm(range(self.max_iter)):
            # 进化状态评估
            if self.archive:
                self.evolutionary_state_evaluation()
            
            # 自适应参数调整
            if self.ES == 'EXPLOITATION':
                self.omega = 0.4
            else:
                self.omega = 0.9
            
            # 选择全局最优（简单随机选择）
            gbest = None
            if self.archive:
                gbest = self.archive[np.random.randint(len(self.archive))]
            
            new_solutions = []
            for i in range(self.pop_size):
                # 更新速度
                r1 = np.random.rand(*self.velocities[i].shape)
                r2 = np.random.rand(*self.velocities[i].shape)
                velocity = (self.omega * self.velocities[i] +
                           self.c1 * r1 * (self.pbest_pos[i] - self.particles[i]) +
                           self.c2 * r2 * (gbest['position'] - self.particles[i]) if gbest else 0)
                
                # 更新位置
                new_pos = np.clip(self.particles[i] + velocity, 0, None).astype(int)
                new_fit, new_viol = self.evaluate(new_pos)
                new_solutions.append({'position': new_pos, 'fitness': new_fit})
                
                # 更新个体最优
                if (new_viol < self.pbest_viol[i] or 
                   (new_viol == self.pbest_viol[i] and 
                    (dominates(new_fit, self.pbest_fit[i]) or 
                     np.random.rand() < 0.5))):
                    self.pbest_pos[i] = new_pos.copy()
                    self.pbest_fit[i] = new_fit
                    self.pbest_viol[i] = new_viol
                
                self.particles[i] = new_pos
            
            # 更新外部档案
            self.update_archive(new_solutions)
        
        return self.archive

def plot_pareto_front(archive, title="Pareto Front", output_file = "resources/resource_scheduling/pareto_front.jpg"):
    """可视化三维Pareto前沿"""
    fig = plt.figure(figsize=(10, 8))
    ax = fig.add_subplot(111, projection='3d')
    
    # 提取目标值
    f1 = [s['fitness'][0] for s in archive]
    f2 = [s['fitness'][1] for s in archive]
    f3 = [s['fitness'][2] for s in archive]
    
    # 绘制散点图
    scatter = ax.scatter(f1, f2, f3, c=f2, cmap='viridis', s=50, alpha=0.8)
    
    # 设置坐标轴标签
    ax.set_xlabel('Time (f1)', fontsize=12)
    ax.set_ylabel('Unsatisfied (f2)', fontsize=12)
    ax.set_zlabel('No-load (f3)', fontsize=12)
    plt.title(title, fontsize=14)
    
    # 添加颜色条
    cbar = fig.colorbar(scatter, pad=0.1)
    cbar.set_label('Unsatisfied Level', rotation=270, labelpad=15)
    
    # 优化视角
    ax.view_init(elev=30, azim=45)
    plt.tight_layout()
    plt.savefig(output_file)

def plot_pareto_front_2d(archive, title="Pareto Front", output_file = "resources/resource_scheduling/pareto_front.jpg"):
    """可视化二维Pareto前沿"""
    fig = plt.figure(figsize=(10, 8))
    ax = fig.add_subplot(111)
    
    # 提取目标值
    f1 = [s['fitness'][0] for s in archive]
    f2 = [s['fitness'][1] for s in archive]
    f1 = np.round(f1, 2)
    f2 = np.round(f2, 2)    
    
    # 绘制散点图
    scatter = ax.scatter(f1, f2, c=f2, cmap='viridis', s=50, alpha=0.8)
    
    # 设置坐标轴标签
    ax.set_xlabel('Time (f1)', fontsize=12)
    ax.set_ylabel('Unsatisfied (f2)', fontsize=12)
    plt.title(title, fontsize=14)
    
    # 添加颜色条
    cbar = fig.colorbar(scatter, pad=0.1)
    cbar.set_label('Unsatisfied Level', rotation=270, labelpad=15)

    plt.tight_layout()
    plt.savefig(output_file)

def calculate_hypervolume(archive, ref_point=None):
    """计算三维目标空间的超体积指标"""
    # 提取目标值矩阵 (N,3)
    points = np.array([sol['fitness'] for sol in archive])
    
    # 自动计算参考点（各目标最大值 + 5%偏移）
    if ref_point is None:
        ref_point = np.max(points, axis=0) * 1.05
    
    # 归一化到参考点坐标系
    normalized = ref_point - points
    
    # 添加原点构成凸包
    points_with_origin = np.vstack([normalized, [0, 0, 0]])
    
    # 计算三维凸包体积
    hull = ConvexHull(points_with_origin)
    return hull.volume    

#style （0：最快路线，1：最短路线，2：避开高速，3：步行）
name_style = {
    0: "最快路线",
    1: "最短路线",
    2: "避开高速",
    3: "步行"
}
def draw(result, name, style, namex, namey, style_flag, vmin, vmax):
    n, m = result.shape
    xLabel = ['%s%d' %(namex, i+1) for i in range(m)]
    yLabel = ['%s%d' %(namey, i+1) for i in range(n)]
 
    font = matplotlib.font_manager.FontProperties(fname="resources/SourceHanSansSC-Normal.otf") 
    fig, ax = plt.subplots()
    plt.subplots_adjust(top=0.85)  # 调整这个值以预留更多或更少的空间
    im = ax.imshow(result, cmap='YlGn', vmin=vmin, vmax=vmax)  # 使用黄绿渐变色
    cbar = ax.figure.colorbar(im, ax=ax)
    cbar.ax.set_ylabel("%s 数量" %(namex), rotation=-90, va="bottom", fontproperties=font)
    ax.set_xticks(np.arange(len(xLabel)), labels=xLabel, fontproperties=font)
    ax.set_yticks(np.arange(len(yLabel)), labels=yLabel, fontproperties=font)
    for i in range(n):
        for j in range(m):
            ax.text(j, i, f"{result[i][j]}", ha="center", va="center", color="black")
    fig.tight_layout(rect=[0, 0, 1, 0.9])  # 确保图像和标题不重叠
    style_n = name_style[style]
    if style_flag:
        plt.title(name + " " + style_n, fontproperties=font, y=1)
    else:
        plt.title(name, fontproperties=font, y=1)
    plt.savefig("resources/" + name + "_" + style_n + ".jpg")
 
if __name__ == "__main__":
    config = read_json("path_planning/tianditu/config.jsonc")
    data_output_path = config["data_output_path"]

    file_pareto_front = "resources/pareto_front_7_1_6"

    with open("data/generate_data_7_1_6.pkl", "rb") as haddle:
        I,J,K,c,d,mu,C = pickle.load(haddle)   

    for style in [0, 1]:
        print ("style", style)
        with open(data_output_path+"_style_%d.pkl" %(style), "rb") as haddle:
            names, destination_name, origin_list, destination, coordinates_list, datas, style = pickle.load(haddle)  
        t = []      
        for distance, duration, mapinfo, routelatlon in datas:
            t.append([float(duration)])
        t = np.array(t)
        print ("t", t)
        start = time.time()
        pso = AMOPSO_ESE(I, J, K, c, d, t, mu, C, pop_size=100, max_iter=2000, archive_size=100)
        archive = pso.run()
        print ("Total Time :", time.time() - start)
        print(f"找到Pareto解数量: {len(archive)}")
        print("最后5个解的目标值：")
        for sol in archive[-5:]:
            print([f"{x:.2f}" for x in sol['fitness']])
        
        for rr in list(range(len(archive)))[-1:]:
            print("第%d个解的情况：" %(rr+1))
            for ii in range(I):
                for jj in range(J):
                    print ("第%d个物资点到第%d个受灾点分配物质情况：" %(ii+1, jj+1), archive[rr]['position'][ii][jj])
        
        # 可视化
        # plot_pareto_front(archive, output_file = file_pareto_front + "_style_%d.jpg" %style)
        vmin = 0
        vmax = 100
        plot_pareto_front_2d(archive, output_file = file_pareto_front + "_style_%d.jpg" %style)
        draw(archive[-1]['position'][:,-1,:], "医院物资分配情况", style, "物资", "医院", True, 0, 700)
        draw(d, "受灾点物资需求情况", style, "物资", "受灾点", False, 400, 900)

        with open("data/resource_scheduling_style%d.pkl" %(style), "wb") as haddle:
            pickle.dump((archive[-1]['position'][:,-1,:], d, t, archive[-1]['fitness'], style), haddle)
        